import json
import os
import subprocess
import threading
import uuid
from datetime import datetime
import shutil
from typing import Optional

import psutil

from fastapi import APIRouter, Depends, UploadFile, File, Query
from sqlalchemy.orm import Session
from database import get_db, SessionLocal
from models.perf import PerfMessages, PerfTask, PerfTaskStatsJson, PerfTaskStatsCsv
from schemas.pref import StartScriptRequest
from models.user import User
from setting import SCRIPT_DIR, REPORT_DIR
from utils.locust_lock import LocustLock
from utils.locust_utils import parse_locust_exceptions_csv
from utils.log_utils import logger
from utils.response_utils import make_response
from utils.token_utils import get_current_user

router = APIRouter(
    prefix="/api/pref/v1",  # 所有路由的前缀
    tags=["性能管理模块"],  # 文档中的标签
    responses={404: {"description": "Not found"}}  # 统一响应
)

task_status_table = {}  # task_id -> status/info
locust_lock = LocustLock()


def finalize_task(task_id: str, exit_code: int):
    """
    压测任务完成后的收尾处理：
    - 更新任务状态
    - 从 JSON 报告文件解析数据写入数据库
    - 释放锁
    """
    logger.info(f"[{task_id}] 正在收尾 finalize_task，exit_code = {exit_code}")
    try:
        with SessionLocal() as db:
            # 获取任务
            task = db.query(PerfTask).filter_by(task_id=task_id).first()
            memory_task = task_status_table.get(task_id)

            if not task:
                logger.warning(f"[{task_id}] 数据库中未找到任务，跳过 finalize")
                return

            # 更新任务状态
            task.status = "finished" if exit_code == 0 else "failed"
            task.end_time = datetime.now()

            # 解析 JSON 报告文件
            json_file = os.path.join(task.path, f"{task.filename}_{task_id}.json")
            failures_csv_file = os.path.join(task.path, f"{task.filename}_{task_id}_failures.csv")

            # 解析json文件
            if not os.path.exists(json_file):
                logger.warning(f"[{task_id}] JSON报告文件不存在: {json_file}")
            else:
                try:
                    with open(json_file, "r", encoding="utf-8") as f:
                        json_data = json.load(f)
                    perf_task_stats_json_db_all_list = []
                    for item in json_data:
                        stat = PerfTaskStatsJson(
                            task_id=task_id,
                            name=item["name"],
                            method=item["method"],
                            num_requests=item["num_requests"],
                            num_failures=item["num_failures"],
                            total_response_time=item["total_response_time"],
                            min_response_time=item["min_response_time"],
                            max_response_time=item["max_response_time"],
                            avg_response_time=(
                                item["total_response_time"] / item["num_requests"]
                                if item["num_requests"] > 0 else 0
                            ),
                            total_content_length=item["total_content_length"],
                            response_times=json.dumps(item.get("response_times", {})),
                            num_reqs_per_sec=json.dumps(item.get("num_reqs_per_sec", {})),
                            num_fail_per_sec=json.dumps(item.get("num_fail_per_sec", {})),
                            tps=(item["num_requests"] / (item["total_response_time"] / 1000)
                                 if item["total_response_time"] > 0 else 0),

                        )
                        perf_task_stats_json_db_all_list.append(stat)
                    db.add_all(perf_task_stats_json_db_all_list)
                    logger.info(f"[{task_id}] 已成功入库接口统计数据")
                except Exception as e:
                    logger.error(f"[{task_id}] JSON报告解析失败: {e}")
            # 解析failures_csv文件
            if not os.path.exists(failures_csv_file):
                logger.warning(f"[{task_id}] failures_csv_file报告文件不存在: {failures_csv_file}")
            else:
                try:
                    results = parse_locust_exceptions_csv(failures_csv_file, encoding="gbk")
                    perf_task_stats_csv_db_all_list = []

                    for result in results:
                        csv_stat = PerfTaskStatsCsv(
                            task_id=task_id,
                            name=result["Name"],
                            method=result["Method"],
                            error=result["Error"],
                            occurrences=result["Occurrences"]
                        )
                        perf_task_stats_csv_db_all_list.append(csv_stat)
                    if perf_task_stats_csv_db_all_list:
                        db.add_all(perf_task_stats_csv_db_all_list)
                    else:
                        logger.info(f"[{task_id}] failures报告中无异常数据")

                except Exception as e:
                    logger.error(f"[{task_id}] failures_csv_file报告解析失败: {e}")

            db.commit()
    except Exception as e:
        logger.error(f"[{task_id}] finalize_task 执行异常: {e}")
    finally:
        locust_lock.release()
        if memory_task:
            memory_task["end_time"] = datetime.now()
            memory_task["status"] = "finished" if exit_code == 0 else "failed"
        logger.info(f"[{task_id}] 锁已释放，收尾完成")


# ===== 后台监控线程 =====

def monitor_process_and_release(process: subprocess.Popen, task_id: str):
    """
    后台线程监控 locust 执行状态
    - 任务执行完后自动释放锁
    - 可扩展记录任务状态到数据库
    """
    logger.info(f"[{task_id}] 后台监控线程启动")

    process.wait()
    exit_code = process.returncode

    finalize_task(task_id, exit_code)


@router.post("/upload_script", summary="上传locust脚本", response_model=dict)
def upload_script(
        file: UploadFile = File(...),
        db: Session = Depends(get_db),
        current_user: User = Depends(get_current_user),
):
    # 1. 验证文件类型
    if not file.filename.endswith(".py"):
        return make_response(400, "只支持py格式")
    filename = file.filename.rsplit(".", 1)[0].strip()
    check = db.query(PerfMessages).filter(PerfMessages.filename == filename).first()  # type: ignore
    if check:
        return make_response(400, "已有重复脚本")

    file_path = SCRIPT_DIR / f"{filename}.py"

    file_path.parent.mkdir(parents=True, exist_ok=True)  # 递归创建目录，避免目录不存在报错

    with open(file_path, "wb") as f:
        f.write(file.file.read())

    db_user = PerfMessages(
        user_id=current_user.id,
        filename=filename,
    )
    db.add(db_user)
    db.commit()
    return make_response(200, "上传成功")


# ===== 启动任务接口 =====

@router.post("/start_script", summary="启动locust脚本", response_model=dict)
def start_script(
        request: StartScriptRequest,
        db: Session = Depends(get_db),
        current_user: User = Depends(get_current_user),
):
    # 1. 判断是否已有任务运行中
    if locust_lock.is_locked():
        return make_response(400, "已有任务在运行，请稍后再试")

    # 2. 查数据库：用户是否上传过该脚本
    script_record = db.query(PerfMessages).filter_by(filename=request.filename, user_id=current_user.id).first()
    if not script_record:
        return make_response(404, "未找到指定脚本")

    # 3. 构造脚本路径
    script_path = SCRIPT_DIR / f"{request.filename}.py"
    if not script_path.exists():
        return make_response(500, "脚本文件不存在")

    # 4. 构造任务ID、日志目录、报告路径
    task_id = str(uuid.uuid4())
    task_dir = REPORT_DIR / task_id
    task_dir.mkdir(parents=True, exist_ok=True)

    log_file_path = task_dir / f"{request.filename}_{task_id}.log"
    report_prefix = task_dir / f"{request.filename}_{task_id}"

    # 5. 构造 locust 命令
    command = [
        "locust",
        "-f", str(script_path),
        "--headless",
        "-u", str(request.users),
        "-r", str(request.spawn_rate),
        "--run-time", request.run_time,
        "--host", request.host,
        "--json-file", str(report_prefix),
        "--csv", str(report_prefix),
        "--csv-full-history",
        "--loglevel", "DEBUG",
        "--logfile", str(log_file_path),
    ]

    try:

        # 7. 启动 locust 压测任务（后台运行，日志写入文件）
        with open(log_file_path, "w") as log_file:
            process = subprocess.Popen(command, stdout=log_file, stderr=log_file, text=True)

        # 6. 创建锁文件，防止并发执行
        locust_lock.acquire(process.pid)  # ✅记录 PID

        # 8. 启动后台线程，监控进程是否结束（释放锁）
        threading.Thread(
            target=monitor_process_and_release,
            args=(process, task_id),
            daemon=True
        ).start()

        # 记录任务状态

        task_status_table[task_id] = {
            "status": "running",
            "start_time": datetime.now(),
            "end_time": None,
            "log_file": str(log_file_path),
        }

        db_task = PerfTask(
            task_id=task_id,
            user_id=current_user.id,
            script_id=script_record.id,
            filename=request.filename,
            status="running",
            users=request.users,
            spawn_rate=request.spawn_rate,
            run_time=request.run_time,
            host=request.host,
            csv_report_path=str(report_prefix) + ".csv",
            log_file_path=str(log_file_path),
            path=task_dir
        )
        db.add(db_task)
        db.commit()

        return make_response(200, "脚本已启动，正在后台压测中", {
            "task_id": task_id,
            "filename": request.filename,
            "bash": " ".join(command),
            "log_file": str(log_file_path)
        })

    except Exception as e:
        locust_lock.release()
        return make_response(500, f"启动失败: {str(e)}")


@router.get("/task_status/{task_id}", summary="查询任务状态")
def task_status(task_id: str, db: Session = Depends(get_db), current_user: User = Depends(get_current_user)):
    task_info = task_status_table.get(task_id)
    if task_info:
        # 若在内存中有，则同步数据库状态
        task = db.query(PerfTask).filter_by(task_id=task_id, user_id=current_user.id).first()
        if task:
            task.status = task_info["status"]
            if task_info["end_time"]:
                task.end_time = task_info["end_time"]
            db.commit()

        return make_response(200, "查询成功", {
            "status": task_info["status"],
            "start_time": task_info["start_time"].isoformat(),
            "end_time": task_info["end_time"].isoformat() if task_info["end_time"] else None,
            "log_file": task_info["log_file"],
        })

    if not task_info:
        # 尝试从数据库查找任务状态（防止服务重启 task_status_table丢失）
        task = db.query(PerfTask).filter_by(task_id=task_id, user_id=current_user.id).first()
        if not task:
            return make_response(404, "任务不存在")
        return make_response(200, "查询成功", {
            "status": task.status,
            "start_time": task.start_time.isoformat() if task.start_time else None,
            "end_time": task.end_time.isoformat() if task.end_time else None,
            "log_file": task.log_file_path,
        })


@router.delete("/delete_script/{filename}", summary="删除locust脚本及历史数据", response_model=dict)
def delete_script(
        filename: str,
        db: Session = Depends(get_db),
        current_user: User = Depends(get_current_user),
):
    record = db.query(PerfMessages).filter_by(filename=filename, user_id=current_user.id).first()
    if not record:
        return make_response(404, "未找到该脚本")

    # 先删除关联的历史任务及相关文件
    tasks = db.query(PerfTask).filter_by(script_id=record.id).all()
    for task in tasks:
        try:
            shutil.rmtree(task.path)
        except Exception as e:
            logger.warning(f"删除任务目录失败: {task.path}，错误: {e}")
        db.delete(task)

    # 删除脚本记录
    db.delete(record)
    db.commit()

    # 删除脚本文件
    script_path = SCRIPT_DIR / f"{filename}.py"
    if script_path.exists():
        script_path.unlink()

    return make_response(200, "脚本及相关历史数据已删除成功")


@router.get("/task_report/{task_id}", summary="查询任务压测报告（含统计和失败信息）", response_model=dict)
def get_task_report(
        task_id: str,
        db: Session = Depends(get_db),
        _=Depends(get_current_user)
):
    task = db.query(PerfTask).filter_by(task_id=task_id).first()
    if not task:
        return make_response(404, "未找到该任务")

    # 接口统计数据（JSON）
    stats_json = []
    for stat in task.task_stats:
        stats_json.append({
            "name": stat.name,
            "method": stat.method,
            "num_requests": stat.num_requests,
            "num_failures": stat.num_failures,
            "total_response_time": stat.total_response_time,
            "min_response_time": stat.min_response_time,
            "max_response_time": stat.max_response_time,
            "avg_response_time": stat.avg_response_time,
            "total_content_length": stat.total_content_length,
            "response_times": json.loads(stat.response_times or "{}"),
            "num_reqs_per_sec": json.loads(stat.num_reqs_per_sec or "{}"),
            "num_fail_per_sec": json.loads(stat.num_fail_per_sec or "{}"),
            "tps": stat.tps
        })
        # 统计所有接口的total_tps
    total_requests = sum(item["num_requests"] for item in stats_json)
    total_time_ms = sum(item["total_response_time"] for item in stats_json)
    total_tps = total_requests / (total_time_ms / 1000) if total_time_ms > 0 else 0

    # 失败统计数据（CSV）
    stats_csv = []
    for fail in task.csv_failures:
        stats_csv.append({
            "name": fail.name,
            "method": fail.method,
            "error": fail.error,
            "occurrences": fail.occurrences
        })

    return make_response(200, "查询成功", {
        "task": {
            "task_id": task.task_id,
            "status": task.status,
            "total_tps": total_tps,
            "start_time": task.start_time.isoformat() if task.start_time else None,
            "end_time": task.end_time.isoformat() if task.end_time else None,
            "users": task.users,
            "spawn_rate": task.spawn_rate,
            "run_time": task.run_time,
            "host": task.host,
            "csv_report_path": task.csv_report_path,
            "log_file_path": task.log_file_path,
        },
        "interface_stats": stats_json,
        "failures": stats_csv,
    })


@router.get("/task_list", summary="查询任务列表（分页+脚本名筛选）", response_model=dict)
def get_task_list(
        page: int = Query(1, ge=1, description="页码，从 1 开始"),
        page_size: int = Query(10, ge=1, le=100, description="每页返回的任务数量"),
        filename: str = Query(None, description="按脚本文件名模糊查询"),
        db: Session = Depends(get_db),
        current_user: User = Depends(get_current_user),
):
    """
    查询当前登录用户的 Locust 压测任务列表。
    支持分页和脚本名模糊筛选。

    - 根据 `user_id` 限定为当前用户上传的任务
    - 支持按 `filename` 模糊过滤（`ilike`）
    - 支持分页：`page`, `page_size`
    - 会自动计算每个任务的总 TPS（吞吐量）

    返回内容包括任务基本信息、压测参数、起止时间、状态、TPS 等。
    """
    # 1. 构造查询条件
    query = db.query(PerfTask).filter(PerfTask.user_id == current_user.id)

    # 2. 模糊筛选脚本文件名（可选）
    if filename:
        query = query.filter(PerfTask.filename.ilike(f"%{filename}%"))

    # 3. 总数量
    total = query.count()

    # 4. 查询当前页数据
    tasks = (
        query.order_by(PerfTask.start_time.desc())
        .offset((page - 1) * page_size)
        .limit(page_size)
        .all()
    )

    # 5. 构建响应数据
    task_list = []
    for task in tasks:
        # 计算该任务的总 TPS
        total_requests = sum([s.num_requests for s in task.task_stats])
        total_time_ms = sum([s.total_response_time for s in task.task_stats])
        total_tps = total_requests / (total_time_ms / 1000) if total_time_ms > 0 else 0

        task_list.append({
            "task_id": task.task_id,  # 任务唯一 ID
            "filename": task.filename,  # 脚本文件名
            "status": task.status,  # 状态：running / finished / failed
            "start_time": task.start_time.isoformat() if task.start_time else None,
            "end_time": task.end_time.isoformat() if task.end_time else None,
            "users": task.users,  # 用户数
            "spawn_rate": task.spawn_rate,  # 启动速率
            "run_time": task.run_time,  # 压测时长
            "host": task.host,  # 被压测地址
            "total_tps": total_tps  # 任务整体 TPS
        })

    # 6. 返回结构化分页数据
    return make_response(200, "查询成功", {
        "total": total,  # 总任务数
        "items": task_list  # 当前页任务列表
    })


@router.get("/script_list", summary="查询当前用户上传的脚本列表", response_model=dict)
def get_script_list(
        page: int = Query(1, ge=1, description="页码"),
        page_size: int = Query(10, ge=1, le=100, description="每页数量"),
        filename: Optional[str] = Query(None, description="按脚本名模糊过滤"),
        db: Session = Depends(get_db),
        current_user: User = Depends(get_current_user),
):
    query = db.query(PerfMessages).filter_by(user_id=current_user.id)

    if filename:
        query = query.filter(PerfMessages.filename.ilike(f"%{filename}%"))

    total = query.count()
    scripts = query.order_by(PerfMessages.created_at.desc()) \
        .offset((page - 1) * page_size) \
        .limit(page_size) \
        .all()

    result = [{
        "id": s.id,
        "filename": s.filename,
        "remarks": s.remarks,
        "created_at": s.created_at.isoformat() if s.created_at else None
    } for s in scripts]

    return make_response(200, "查询成功", {
        "total": total,
        "items": result
    })
